import multiprocessing
import os
import re
import sys
from functools import partial

import mwparserfromhell
from tqdm import tqdm
from zkl_promptui import confirm_clear_dir_path

from zkl_datasets import ChunkedBinaryDatasetWriter, load_dataset

root_dir_path = os.path.join(os.path.dirname(__file__), "../../..")
sys.path.append(root_dir_path)

from scripts.config import datasets_dir_path

# Source: https://en.wikipedia.org/wiki/List_of_Wikipedias (accessed 3/1/2019)
# Removed because no articles: hz.
LANGUAGES = [
    "aa",
    "ab",
    "ace",
    "ady",
    "af",
    "ak",
    "als",
    "am",
    "an",
    "ang",
    "ar",
    "arc",
    "arz",
    "as",
    "ast",
    "atj",
    "av",
    "ay",
    "az",
    "azb",
    "ba",
    "bar",
    "bat-smg",
    "bcl",
    "be",
    "be-x-old",
    "bg",
    "bh",
    "bi",
    "bjn",
    "bm",
    "bn",
    "bo",
    "bpy",
    "br",
    "bs",
    "bug",
    "bxr",
    "ca",
    "cbk-zam",
    "cdo",
    "ce",
    "ceb",
    "ch",
    "cho",
    "chr",
    "chy",
    "ckb",
    "co",
    "cr",
    "crh",
    "cs",
    "csb",
    "cu",
    "cv",
    "cy",
    "da",
    "de",
    "din",
    "diq",
    "dsb",
    "dty",
    "dv",
    "dz",
    "ee",
    "el",
    "eml",
    "en",
    "eo",
    "es",
    "et",
    "eu",
    "ext",
    "fa",
    "ff",
    "fi",
    "fiu-vro",
    "fj",
    "fo",
    "fr",
    "frp",
    "frr",
    "fur",
    "fy",
    "ga",
    "gag",
    "gan",
    "gd",
    "gl",
    "glk",
    "gn",
    "gom",
    "gor",
    "got",
    "gu",
    "gv",
    "ha",
    "hak",
    "haw",
    "he",
    "hi",
    "hif",
    "ho",
    "hr",
    "hsb",
    "ht",
    "hu",
    "hy",
    "ia",
    "id",
    "ie",
    "ig",
    "ii",
    "ik",
    "ilo",
    "inh",
    "io",
    "is",
    "it",
    "iu",
    "ja",
    "jam",
    "jbo",
    "jv",
    "ka",
    "kaa",
    "kab",
    "kbd",
    "kbp",
    "kg",
    "ki",
    "kj",
    "kk",
    "kl",
    "km",
    "kn",
    "ko",
    "koi",
    "krc",
    "ks",
    "ksh",
    "ku",
    "kv",
    "kw",
    "ky",
    "la",
    "lad",
    "lb",
    "lbe",
    "lez",
    "lfn",
    "lg",
    "li",
    "lij",
    "lmo",
    "ln",
    "lo",
    "lrc",
    "lt",
    "ltg",
    "lv",
    "mai",
    "map-bms",
    "mdf",
    "mg",
    "mh",
    "mhr",
    "mi",
    "min",
    "mk",
    "ml",
    "mn",
    "mr",
    "mrj",
    "ms",
    "mt",
    "mus",
    "mwl",
    "my",
    "myv",
    "mzn",
    "na",
    "nah",
    "nap",
    "nds",
    "nds-nl",
    "ne",
    "new",
    "ng",
    "nl",
    "nn",
    "no",
    "nov",
    "nrm",
    "nso",
    "nv",
    "ny",
    "oc",
    "olo",
    "om",
    "or",
    "os",
    "pa",
    "pag",
    "pam",
    "pap",
    "pcd",
    "pdc",
    "pfl",
    "pi",
    "pih",
    "pl",
    "pms",
    "pnb",
    "pnt",
    "ps",
    "pt",
    "qu",
    "rm",
    "rmy",
    "rn",
    "ro",
    "roa-rup",
    "roa-tara",
    "ru",
    "rue",
    "rw",
    "sa",
    "sah",
    "sat",
    "sc",
    "scn",
    "sco",
    "sd",
    "se",
    "sg",
    "sh",
    "si",
    "simple",
    "sk",
    "sl",
    "sm",
    "sn",
    "so",
    "sq",
    "sr",
    "srn",
    "ss",
    "st",
    "stq",
    "su",
    "sv",
    "sw",
    "szl",
    "ta",
    "tcy",
    "te",
    "tet",
    "tg",
    "th",
    "ti",
    "tk",
    "tl",
    "tn",
    "to",
    "tpi",
    "tr",
    "ts",
    "tt",
    "tum",
    "tw",
    "ty",
    "tyv",
    "udm",
    "ug",
    "uk",
    "ur",
    "uz",
    "ve",
    "vec",
    "vep",
    "vi",
    "vls",
    "vo",
    "wa",
    "war",
    "wo",
    "wuu",
    "xal",
    "xh",
    "xmf",
    "yi",
    "yo",
    "za",
    "zea",
    "zh",
    "zh-classical",
    "zh-min-nan",
    "zh-yue",
    "zu",
]

# Source: for each Wikipedia language code (example shown for "ab"), aliases for namespaces -2 and 6 accessed via this API call:
# https://ab.wikipedia.org/w/api.php?action=query&meta=siteinfo&siprop=namespacealiases|namespaces&format=json&formatversion=2 (accessed 12/21/2021)
ALIASES_MEDIA = {
    "ab":           ["Медиа", "Файл", "Афаил", "Амедиа", "Изображение"],
    "ace":          ["Beureukaih", "Gambar", "Alat", "Berkas"],
    "ady":          ["Медиа"],
    "af":           ["Lêer", "Beeld"],
    "als":          ["Medium", "Datei", "Bild"],
    "am":           ["ፋይል", "ስዕል"],
    "an":           ["Imachen", "Imagen"],
    "ang":          ["Ymele", "Biliþ"],
    "ar":           ["ميديا", "صورة", "وسائط", "ملف"],
    "arc":          ["ܠܦܦܐ", "ܡܝܕܝܐ"],
    "arz":          ["ميديا", "صورة", "وسائط", "ملف"],
    "as":           ["চিত্ৰ", "चित्र", "চিত্র", "মাধ্যম"],
    "ast":          ["Imaxen", "Ficheru", "Imaxe", "Archivu", "Imagen", "Medios"],
    "atj":          ["Tipatcimoctakewin", "Natisinahikaniwoc"],
    "av":           ["Медиа", "Файл", "Изображение"],
    "ay":           ["Medio", "Archivo", "Imagen"],
    "az":           ["Mediya", "Şəkil", "Fayl"],
    "azb":          ["رسانه", "تصویر", "مدیا", "فایل", "رسانه‌ای"],
    "ba":           ["Медиа", "Рәсем", "Файл", "Изображение"],
    "bar":          ["Medium", "Datei", "Bild"],
    "bat-smg":      ["Vaizdas", "Medėjė", "Abruozdielis"],
    "bcl":          ["Medio", "Ladawan"],
    "be":           ["Мультымедыя", "Файл", "Выява"],
    "be-x-old":     ["Мэдыя", "Файл", "Выява"],
    "bg":           ["Медия", "Файл", "Картинка"],
    "bh":           ["मीडिया", "चित्र"],
    "bjn":          ["Barakas", "Gambar", "Berkas"],
    "bm":           ["Média", "Fichier"],
    "bn":           ["চিত্র", "মিডিয়া"],
    "bpy":          ["ছবি", "মিডিয়া"],
    "br":           ["Skeudenn", "Restr"],
    "bs":           ["Mediji", "Slika", "Datoteka", "Medija"],
    "bug":          ["Gambar", "Berkas"],
    "bxr":          ["Файл", "Меди", "Изображение"],
    "ca":           ["Fitxer", "Imatge"],
    "cbk-zam":      ["Medio", "Archivo", "Imagen"],
    "cdo":          ["文件", "媒體", "圖像", "檔案"],
    "ce":           ["Хlум", "Медиа", "Сурт", "Файл", "Медйа", "Изображение"],
    "ceb":          ["Payl", "Medya", "Imahen"],
    "ch":           ["Litratu"],
    "ckb":          ["میدیا", "پەڕگە"],
    "co":           ["Immagine"],
    "crh":          ["Медиа", "Resim", "Файл", "Fayl", "Ресим"],
    "cs":           ["Soubor", "Média", "Obrázok"],
    "csb":          ["Òbrôzk", "Grafika"],
    "cu":           ["Видъ", "Ви́дъ", "Дѣло", "Срѣдьства"],
    "cv":           ["Медиа", "Ӳкерчĕк", "Изображение"],
    "cy":           ["Delwedd"],
    "da":           ["Billede", "Fil"],
    "de":           ["Medium", "Datei", "Bild"],
    "din":          ["Ciɛl", "Apamduööt"],
    "diq":          ["Medya", "Dosya"],
    "dsb":          ["Wobraz", "Dataja", "Bild", "Medija"],
    "dty":          ["चित्र", "मिडिया"],
    "dv":           ["ފައިލު", "މީޑިއާ", "ފައިލް"],
    "el":           ["Εικόνα", "Αρχείο", "Μέσο", "Μέσον"],
    "eml":          ["Immagine"],
    "eo":           ["Dosiero", "Aŭdvidaĵo"],
    "es":           ["Medio", "Archivo", "Imagen"],
    "et":           ["Pilt", "Fail", "Meedia"],
    "eu":           ["Irudi", "Fitxategi"],
    "ext":          ["Archivu", "Imagen", "Mediu"],
    "fa":           ["رسانه", "تصویر", "مدیا", "پرونده", "رسانه‌ای"],
    "ff":           ["Média", "Fichier"],
    "fi":           ["Kuva", "Tiedosto"],
    "fiu-vro":      ["Pilt", "Meediä"],
    "fo":           ["Miðil", "Mynd"],
    "fr":           ["Média", "Fichier"],
    "frp":          ["Émâge", "Fichiér", "Mèdia"],
    "frr":          ["Medium", "Datei", "Bild"],
    "fur":          ["Immagine", "Figure"],
    "fy":           ["Ofbyld"],
    "ga":           ["Íomhá", "Meán"],
    "gag":          ["Mediya", "Medya", "Resim", "Dosya", "Dosye"],
    "gan":          ["媒体文件", "文件", "文檔", "档案", "媒體", "图像", "圖像", "媒体", "檔案"],
    "gd":           ["Faidhle", "Meadhan"],
    "gl":           ["Imaxe", "Ficheiro", "Arquivo", "Imagem"],
    "glk":          ["رسانه", "تصویر", "پرونده", "فاىل", "رسانه‌ای", "مديا"],
    "gn":           ["Medio", "Imagen", "Ta'ãnga"],
    "gom":          ["माध्यम", "मिडिया", "फायल"],
    "gor":          ["Gambar", "Berkas"],
    "got":          ["𐍆𐌴𐌹𐌻𐌰"],
    "gu":           ["દ્રશ્ય-શ્રાવ્ય (મિડિયા)", "દ્રશ્ય-શ્રાવ્ય_(મિડિયા)", "ચિત્ર"],
    "gv":           ["Coadan", "Meanyn"],
    "hak":          ["文件", "媒體", "圖像", "檔案"],
    "haw":          ["Kiʻi", "Waihona", "Pāpaho"],
    "he":           ["תמונה", "קו", "מדיה", "קובץ"],
    "hi":           ["मीडिया", "चित्र"],
    "hif":          ["file", "saadhan"],
    "hr":           ["Mediji", "DT", "Slika", "F", "Datoteka"],
    "hsb":          ["Wobraz", "Dataja", "Bild"],
    "ht":           ["Imaj", "Fichye", "Medya"],
    "hu":           ["Kép", "Fájl", "Média"],
    "hy":           ["Պատկեր", "Մեդիա"],
    "ia":           ["Imagine", "Multimedia"],
    "id":           ["Gambar", "Berkas"],
    "ig":           ["Nká", "Midia", "Usòrò", "Ákwúkwó orünotu", "Ákwúkwó_orünotu"],
    "ii":           ["媒体文件", "文件", "档案", "图像", "媒体"],
    "ilo":          ["Midia", "Papeles"],
    "inh":          ["Медиа", "Файл", "Изображение"],
    "io":           ["Imajo", "Arkivo"],
    "is":           ["Miðill", "Mynd"],
    "it":           ["Immagine"],
    "ja":           ["メディア", "ファイル", "画像"],
    "jbo":          ["velsku", "datnyvei"],
    "jv":           ["Barkas", "Medhia", "Gambar", "Médhia"],
    "ka":           ["მედია", "სურათი", "ფაილი"],
    "kaa":          ["Swret", "Таспа", "سۋرەت", "Taspa", "Su'wret", "Сурет", "تاسپا"],
    "kab":          ["Tugna"],
    "kbd":          ["Медиа", "Файл"],
    "kbp":          ["Média", "Fichier"],
    "kg":           ["Fisye"],
    "kk":           ["Swret", "سۋرەت", "Таспа", "Taspa", "Сурет", "تاسپا"],
    "kl":           ["Billede", "Fiileq", "Fil"],
    "km":           ["ឯកសារ", "រូបភាព", "មេឌា", "មីឌា"],
    "kn":           ["ಚಿತ್ರ", "ಮೀಡಿಯ"],
    "ko":           ["미디어", "파일", "그림"],
    "koi":          ["Медиа", "Файл", "Изображение"],
    "krc":          ["Медиа", "Файл", "Изображение"],
    "ks":           ["میڈیا", "فَیِل"],
    "ksh":          ["Beld", "Meedije", "Medie", "Belld", "Medium", "Datei", "Meedijum", "Bild"],
    "ku":           ["میدیا", "پەڕگە", "Medya", "Wêne"],
    "kv":           ["Медиа", "Файл", "Изображение"],
    "kw":           ["Restren"],
    "ky":           ["Медиа", "Файл"],
    "la":           ["Imago", "Fasciculus"],
    "lad":          ["Dossia", "Medya", "Archivo", "Dosya", "Imagen", "Meddia"],
    "lb":           ["Fichier", "Bild"],
    "lbe":          ["Медиа", "Сурат", "Изображение"],
    "lez":          ["Медиа", "Mediya", "Файл", "Şəkil", "Изображение"],
    "lfn":          ["Fix"],
    "li":           ["Afbeelding", "Plaetje", "Aafbeilding"],
    "lij":          ["Immaggine", "Immagine"],
    "lmo":          ["Immagine", "Imàjine", "Archivi"],
    "ln":           ["Média", "Fichier"],
    "lo":           ["ສື່ອ", "ສື່", "ຮູບ"],
    "lrc":          ["رسانه", "تصویر", "رسانه‌ای", "جانیا", "أسگ", "ڤارئسگأر"],
    "lt":           ["Vaizdas", "Medija"],
    "ltg":          ["Medeja", "Fails"],
    "lv":           ["Attēls"],
    "mai":          ["मेडिया", "फाइल"],
    "map-bms":      ["Barkas", "Medhia", "Gambar", "Médhia"],
    "mdf":          ["Медиа", "Няйф", "Изображение"],
    "mg":           ["Rakitra", "Sary", "Média"],
    "mhr":          ["Медиа", "Файл", "Изображение"],
    "min":          ["Gambar", "Berkas"],
    "mk":           ["Податотека", "Медија", "Медиум", "Слика"],
    "ml":           ["പ്രമാണം", "ചി", "മീഡിയ", "പ്ര", "ചിത്രം"],
    "mn":           ["Медиа", "Файл", "Зураг"],
    "mr":           ["चित्र", "मिडिया"],
    "mrj":          ["Медиа", "Файл", "Изображение"],
    "ms":           ["Fail", "Imej"],
    "mt":           ["Midja", "Medja", "Stampa"],
    "mwl":          ["Multimédia", "Fexeiro", "Ficheiro", "Arquivo", "Imagem"],
    "my":           ["ဖိုင်", "မီဒီယာ"],
    "myv":          ["Медия", "Артовкс", "Изображение"],
    "mzn":          ["رسانه", "تصویر", "مه‌دیا", "مدیا", "پرونده", "رسانه‌ای"],
    "nah":          ["Mēdiatl", "Īxiptli", "Imagen"],
    "nap":          ["Fiùra", "Immagine"],
    "nds":          ["Datei", "Bild"],
    "nds-nl":       ["Ofbeelding", "Afbeelding", "Bestaand"],
    "ne":           ["मीडिया", "चित्र"],
    "new":          ["किपा", "माध्यम"],
    "nl":           ["Bestand", "Afbeelding"],
    "nn":           ["Fil", "Bilde", "Filpeikar"],
    "no":           ["Fil", "Medium", "Bilde"],
    "nov":          [],
    "nrm":          ["Média", "Fichier"],
    "nso":          ["Seswantšho"],
    "nv":           ["Eʼelyaaígíí"],
    "oc":           ["Imatge", "Fichièr", "Mèdia"],
    "olo":          ["Kuva", "Medii", "Failu"],
    "or":           ["ମାଧ୍ୟମ", "ଫାଇଲ"],
    "os":           ["Ныв", "Медиа", "Файл", "Изображение"],
    "pa":           ["ਤਸਵੀਰ", "ਮੀਡੀਆ"],
    "pcd":          ["Média", "Fichier"],
    "pdc":          ["Medium", "Datei", "Bild", "Feil"],
    "pfl":          ["Dadai", "Medium", "Datei", "Bild"],
    "pi":           ["मीडिया", "पटिमा"],
    "pl":           ["Plik", "Grafika"],
    "pms":          ["Figura", "Immagine"],
    "pnb":          ["میڈیا", "تصویر", "فائل"],
    "pnt":          ["Εικόνα", "Αρχείον", "Εικόναν", "Μέσον"],
    "ps":           ["انځور", "رسنۍ", "دوتنه"],
    "pt":           ["Multimédia", "Ficheiro", "Arquivo", "Imagem"],
    "qu":           ["Midya", "Imagen", "Rikcha"],
    "rm":           ["Multimedia", "Datoteca"],
    "rmy":          ["Fişier", "Mediya", "Chitro", "Imagine"],
    "ro":           ["Fişier", "Imagine", "Fișier"],
    "roa-rup":      ["Fişier", "Imagine", "Fișier"],
    "roa-tara":     ["Immagine"],
    "ru":           ["Медиа", "Файл", "Изображение"],
    "rue":          ["Медіа", "Медиа", "Файл", "Изображение", "Зображення"],
    "rw":           ["Dosiye", "Itangazamakuru"],
    "sa":           ["चित्रम्", "माध्यमम्", "सञ्चिका", "माध्यम", "चित्रं"],
    "sah":          ["Миэдьийэ", "Ойуу", "Билэ", "Изображение"],
    "sat":          ["ᱨᱮᱫ", "ᱢᱤᱰᱤᱭᱟ"],
    "sc":           ["Immàgini"],
    "scn":          ["Immagine", "Mmàggini", "Mèdia"],
    "sd":           ["عڪس", "ذريعات", "فائل"],
    "se":           ["Fiila"],
    "sg":           ["Média", "Fichier"],
    "sh":           ["Mediji", "Slika", "Медија", "Datoteka", "Medija", "Слика"],
    "si":           ["රූපය", "මාධ්‍යය", "ගොනුව"],
    "sk":           ["Súbor", "Obrázok", "Médiá"],
    "sl":           ["Slika", "Datoteka"],
    "sq":           ["Figura", "Skeda"],
    "sr":           ["Датотека", "Medij", "Slika", "Медија", "Datoteka", "Медиј", "Medija", "Слика"],
    "srn":          ["Afbeelding", "Gefre"],
    "stq":          ["Bielde", "Bild"],
    "su":           ["Média", "Gambar"],
    "sv":           ["Fil", "Bild"],
    "sw":           ["Faili", "Picha"],
    "szl":          ["Plik", "Grafika"],
    "ta":           ["படிமம்", "ஊடகம்"],
    "tcy":          ["ಮಾದ್ಯಮೊ", "ಫೈಲ್"],
    "te":           ["ఫైలు", "దస్త్రం", "బొమ్మ", "మీడియా"],
    "tet":          ["Imajen", "Arquivo", "Imagem"],
    "tg":           ["Акс", "Медиа"],
    "th":           ["ไฟล์", "สื่อ", "ภาพ"],
    "ti":           ["ፋይል", "ሜድያ"],
    "tk":           ["Faýl"],
    "tl":           ["Midya", "Talaksan"],
    "tpi":          ["Fail"],
    "tr":           ["Medya", "Resim", "Dosya", "Ortam"],
    "tt":           ["Медиа", "Рәсем", "Файл", "Räsem", "Изображение"],
    "ty":           ["Média", "Fichier"],
    "tyv":          ["Медиа", "Файл", "Изображение"],
    "udm":          ["Медиа", "Файл", "Суред", "Изображение"],
    "ug":           ["ۋاسىتە", "ھۆججەت"],
    "uk":           ["Медіа", "Медиа", "Файл", "Изображение", "Зображення"],
    "ur":           ["میڈیا", "تصویر", "وسیط", "زریعہ", "فائل", "ملف"],
    "uz":           ["Mediya", "Tasvir", "Fayl"],
    "vec":          ["Immagine", "Imàjine", "Mèdia"],
    "vep":          ["Pilt", "Fail"],
    "vi":           ["Phương_tiện", "Tập_tin", "Hình", "Tập tin", "Phương tiện"],
    "vls":          ["Afbeelding", "Ofbeeldienge"],
    "vo":           ["Ragiv", "Magod", "Nünamakanäd"],
    "wa":           ["Imådje"],
    "war":          ["Medya", "Fayl", "Paypay"],
    "wo":           ["Xibaarukaay", "Dencukaay"],
    "wuu":          ["文件", "档案", "图像", "媒体"],
    "xal":          ["Аһар", "Боомг", "Изображение", "Зург"],
    "xmf":          ["მედია", "სურათი", "ფაილი"],
    "yi":           ["מעדיע", "תמונה", "טעקע", "בילד"],
    "yo":           ["Fáìlì", "Amóhùnmáwòrán", "Àwòrán"],
    "za":           ["媒体文件", "文件", "档案", "图像", "媒体"],
    "zea":          ["Afbeelding", "Plaetje"],
    "zh":           ["媒体文件", "F", "文件", "媒體", "档案", "图像", "圖像", "媒体", "檔案"],
    "zh-classical": ["文件", "媒體", "圖像", "檔案"],
    "zh-min-nan":   ["tóng-àn", "文件", "媒體", "Mûi-thé", "圖像", "檔案"],
    "zh-yue":       ["檔", "档", "文件", "图", "媒體", "圖", "档案", "图像", "圖像", "媒体", "檔案"],
}

# Source: for each Wikipedia language code (example shown for "ab"), aliases for namespace 14 accessed via this API call:
# https://ab.wikipedia.org/w/api.php?action=query&meta=siteinfo&siprop=namespacealiases|namespaces&format=json&formatversion=2 (accessed 12/21/2021)
ALIASES_CATEGORY = {
    "ab":           ["Категория", "Акатегориа"],
    "ace":          ["Kawan", "Kategori"],
    "af":           ["Kategorie"],
    "ak":           ["Nkyekyem"],
    "als":          ["Kategorie"],
    "am":           ["መደብ"],
    "an":           ["Categoría"],
    "ang":          ["Flocc"],
    "ar":           ["تصنيف"],
    "arc":          ["ܣܕܪܐ"],
    "arz":          ["تصنيف"],
    "as":           ["CAT", "শ্ৰেণী", "श्रेणी", "শ্রেণী"],
    "ast":          ["Categoría"],
    "atj":          ["Tipanictawin"],
    "av":           ["Категория"],
    "ay":           ["Categoría"],
    "az":           ["Kateqoriya"],
    "azb":          ["بؤلمه"],
    "ba":           ["Төркөм", "Категория"],
    "bar":          ["Kategorie"],
    "bat-smg":      ["Kategorija", "Kateguorėjė"],
    "bcl":          ["Kategorya"],
    "be":           ["Катэгорыя"],
    "be-x-old":     ["Катэгорыя"],
    "bg":           ["Категория"],
    "bh":           ["श्रेणी"],
    "bjn":          ["Tumbung", "Kategori"],
    "bm":           ["Catégorie"],
    "bn":           ["বিষয়শ্রেণী", "വിഭാഗം"],
    "bpy":          ["থাক"],
    "br":           ["Rummad"],
    "bs":           ["Kategorija"],
    "bug":          ["Kategori"],
    "bxr":          ["Категори", "Категория"],
    "ca":           ["Categoria"],
    "cbk-zam":      ["Categoría"],
    "cdo":          ["分類"],
    "ce":           ["Категори", "Тоба", "Кадегар"],
    "ceb":          ["Kategoriya"],
    "ch":           ["Katigoria"],
    "ckb":          ["پ", "پۆل"],
    "co":           ["Categoria"],
    "crh":          ["Категория", "Kategoriya"],
    "cs":           ["Kategorie"],
    "csb":          ["Kategòrëjô"],
    "cu":           ["Катигорї", "Категория", "Катигорїꙗ"],
    "cv":           ["Категори"],
    "cy":           ["Categori"],
    "da":           ["Kategori"],
    "de":           ["Kategorie"],
    "din":          ["Bekätakthook"],
    "diq":          ["Kategoriye", "Kategori"],
    "dsb":          ["Kategorija"],
    "dty":          ["श्रेणी"],
    "dv":           ["ޤިސްމު"],
    "el":           ["Κατηγορία"],
    "eml":          ["Categoria"],
    "eo":           ["Kategorio"],
    "es":           ["CAT", "Categoría"],
    "et":           ["Kategooria"],
    "eu":           ["Kategoria"],
    "ext":          ["Categoría", "Categoria"],
    "fa":           ["رده"],
    "ff":           ["Catégorie"],
    "fi":           ["Luokka"],
    "fiu-vro":      ["Katõgooria"],
    "fo":           ["Bólkur"],
    "fr":           ["Catégorie"],
    "frp":          ["Catègorie"],
    "frr":          ["Kategorie"],
    "fur":          ["Categorie"],
    "fy":           ["Kategory"],
    "ga":           ["Rang", "Catagóir"],
    "gag":          ["Kategori", "Kategoriya"],
    "gan":          ["分類", "分类"],
    "gd":           ["Roinn-seòrsa"],
    "gl":           ["Categoría"],
    "glk":          ["جرگه", "رده"],
    "gn":           ["Ñemohenda"],
    "gom":          ["वर्ग", "श्रेणी"],
    "gor":          ["Dalala"],
    "got":          ["𐌷𐌰𐌽𐍃𐌰"],
    "gu":           ["શ્રેણી", "CAT", "શ્રે"],
    "gv":           ["Ronney"],
    "hak":          ["分類"],
    "haw":          ["Māhele"],
    "he":           ["קטגוריה", "קט"],
    "hi":           ["श्र", "श्रेणी"],
    "hif":          ["vibhag"],
    "hr":           ["CT", "KT", "Kategorija"],
    "hsb":          ["Kategorija"],
    "ht":           ["Kategori"],
    "hu":           ["Kategória"],
    "hy":           ["Կատեգորիա"],
    "ia":           ["Categoria"],
    "id":           ["Kategori"],
    "ie":           ["Categorie"],
    "ig":           ["Ébéonọr", "Òtù"],
    "ii":           ["分类"],
    "ilo":          ["Kategoria"],
    "inh":          ["ОагӀат"],
    "io":           ["Kategorio"],
    "is":           ["Flokkur"],
    "it":           ["CAT", "Categoria"],
    "ja":           ["カテゴリ"],
    "jbo":          ["klesi"],
    "jv":           ["Kategori"],
    "ka":           ["კატეგორია"],
    "kaa":          ["Sanat", "Kategoriya", "Санат", "سانات"],
    "kab":          ["Taggayt"],
    "kbd":          ["Категория", "Категориэ"],
    "kbp":          ["Catégorie"],
    "kg":           ["Kalasi"],
    "kk":           ["Sanat", "Санат", "سانات"],
    "kl":           ["Sumut_atassuseq", "Kategori", "Sumut atassuseq"],
    "km":           ["ចំនាត់ថ្នាក់ក្រុម", "ចំណាត់ក្រុម", "ចំណាត់ថ្នាក់ក្រុម"],
    "kn":           ["ವರ್ಗ"],
    "ko":           ["분류"],
    "koi":          ["Категория"],
    "krc":          ["Категория"],
    "ks":           ["زٲژ"],
    "ksh":          ["Saachjropp", "Saachjrop", "Katejori", "Kategorie", "Saachjrupp", "Kattejori", "Sachjrop"],
    "ku":           ["Kategorî", "پۆل"],
    "kv":           ["Категория"],
    "kw":           ["Class", "Klass"],
    "ky":           ["Категория"],
    "la":           ["Categoria"],
    "lad":          ["Kateggoría", "Katēggoría", "Categoría"],
    "lb":           ["Kategorie"],
    "lbe":          ["Категория"],
    "lez":          ["Категория"],
    "lfn":          ["Categoria"],
    "li":           ["Categorie", "Kategorie"],
    "lij":          ["Categorîa", "Categoria"],
    "lmo":          ["Categuria", "Categoria"],
    "ln":           ["Catégorie"],
    "lo":           ["ໝວດ"],
    "lrc":          ["دأسە"],
    "lt":           ["Kategorija"],
    "ltg":          ["Kategoreja"],
    "lv":           ["Kategorija"],
    "mai":          ["CA", "श्रेणी"],
    "map-bms":      ["Kategori"],
    "mdf":          ["Категорие", "Категория"],
    "mg":           ["Sokajy", "Catégorie"],
    "mhr":          ["Категория", "Категорий"],
    "min":          ["Kategori"],
    "mk":           ["Категорија"],
    "ml":           ["വിഭാഗം", "വി", "വർഗ്ഗം", "വ"],
    "mn":           ["Ангилал"],
    "mr":           ["वर्ग"],
    "mrj":          ["Категори", "Категория"],
    "ms":           ["Kategori"],
    "mt":           ["Kategorija"],
    "mwl":          ["Catadorie", "Categoria"],
    "my":           ["ကဏ္ဍ"],
    "myv":          ["Категория"],
    "mzn":          ["رج", "رده"],
    "nah":          ["Neneuhcāyōtl", "Categoría"],
    "nap":          ["Categurìa", "Categoria"],
    "nds":          ["Kategorie"],
    "nds-nl":       ["Categorie", "Kattegerie", "Kategorie"],
    "ne":           ["श्रेणी"],
    "new":          ["पुचः"],
    "nl":           ["Categorie"],
    "nn":           ["Kategori"],
    "no":           ["Kategori"],
    "nrm":          ["Catégorie"],
    "nso":          ["Setensele"],
    "nv":           ["Tʼááłáhági_átʼéego", "Tʼááłáhági átʼéego"],
    "oc":           ["Categoria"],
    "olo":          ["Kategourii"],
    "or":           ["ବିଭାଗ", "ଶ୍ରେଣୀ"],
    "os":           ["Категори"],
    "pa":           ["ਸ਼੍ਰੇਣੀ"],
    "pcd":          ["Catégorie"],
    "pdc":          ["Abdeeling", "Kategorie"],
    "pfl":          ["Kadegorie", "Sachgrubb", "Kategorie"],
    "pi":           ["विभाग"],
    "pl":           ["Kategoria"],
    "pms":          ["Categorìa"],
    "pnb":          ["گٹھ"],
    "pnt":          ["Κατηγορίαν"],
    "ps":           ["وېشنيزه"],
    "pt":           ["Categoria"],
    "qu":           ["Katiguriya"],
    "rm":           ["Categoria"],
    "rmy":          ["Shopni"],
    "ro":           ["Categorie"],
    "roa-rup":      ["Categorie"],
    "roa-tara":     ["Categoria"],
    "ru":           ["Категория", "К"],
    "rue":          ["Категория", "Катеґорія"],
    "rw":           ["Ikiciro"],
    "sa":           ["वर्गः"],
    "sah":          ["Категория"],
    "sat":          ["ᱛᱷᱚᱠ"],
    "sc":           ["Categoria"],
    "scn":          ["Catigurìa"],
    "sd":           ["زمرو"],
    "se":           ["Kategoriija"],
    "sg":           ["Catégorie"],
    "sh":           ["Kategorija", "Категорија"],
    "si":           ["ප්‍රවර්ගය"],
    "sk":           ["Kategória"],
    "sl":           ["Kategorija"],
    "sq":           ["Kategoria", "Kategori"],
    "sr":           ["Kategorija", "Категорија"],
    "srn":          ["Categorie", "Guru"],
    "stq":          ["Kategorie"],
    "su":           ["Kategori"],
    "sv":           ["Kategori"],
    "sw":           ["Jamii"],
    "szl":          ["Kategoryjo", "Kategoria"],
    "ta":           ["பகுப்பு"],
    "tcy":          ["ವರ್ಗೊ"],
    "te":           ["వర్గం"],
    "tet":          ["Kategoría", "Kategoria"],
    "tg":           ["Гурӯҳ"],
    "th":           ["หมวดหมู่"],
    "ti":           ["መደብ"],
    "tk":           ["Kategoriýa"],
    "tl":           ["Kategorya", "Kaurian"],
    "tpi":          ["Grup"],
    "tr":           ["Kategori", "KAT"],
    "tt":           ["Төркем", "Törkem", "Категория"],
    "ty":           ["Catégorie"],
    "tyv":          ["Аңгылал", "Категория"],
    "udm":          ["Категория"],
    "ug":           ["تۈر"],
    "uk":           ["Категория", "Категорія"],
    "ur":           ["زمرہ"],
    "uz":           ["Turkum", "Kategoriya"],
    "vec":          ["Categoria"],
    "vep":          ["Kategorii"],
    "vi":           ["Thể_loại", "Thể loại"],
    "vls":          ["Categorie"],
    "vo":           ["Klad"],
    "wa":           ["Categoreye"],
    "war":          ["Kaarangay"],
    "wo":           ["Wàll", "Catégorie"],
    "wuu":          ["分类"],
    "xal":          ["Янз", "Әәшл"],
    "xmf":          ["კატეგორია"],
    "yi":           ["קאטעגאריע", "קאַטעגאָריע"],
    "yo":           ["Ẹ̀ka"],
    "za":           ["分类"],
    "zea":          ["Categorie"],
    "zh":           ["分类", "分類", "CAT"],
    "zh-classical": ["分類", "CAT"],
    "zh-min-nan":   ["分類", "Lūi-pia̍t"],
    "zh-yue":       ["分类", "分類", "类", "類"],
}

re_magic = re.compile("__[A-Z]*__", flags=re.UNICODE)
re_box = re.compile(r"\{\|.*?\|}", flags=re.DOTALL | re.UNICODE)
re_lsw = re.compile(r"-\{([a-z]+?-?[a-z]+?:(.*?))+;?}-", flags=re.DOTALL | re.UNICODE)


def clean_magic(text: str) -> str:
    text = re.sub(re_magic, "", text)
    return text


def clean_box(text: str) -> str:
    text = re.sub(re_box, "", text)
    return text


def clean_lsw(text: str):
    re_lsw = re.compile(r"-\{([a-z-]+):(.*?);?(.*?)}-", flags=re.DOTALL | re.UNICODE)
    text = re.sub(re_lsw, "\2", text)
    return text


def parse_and_clean_wikicode(wikicode_text: str, lang: str) -> str:
    """Strips formatting and unwanted sections from raw page content."""
    wikicode = mwparserfromhell.parse(wikicode_text)

    # Filters for file/image links.
    media_prefixes = "|".join(["File", "Image", "Media"] + ALIASES_MEDIA.get(lang, []))
    re_rm_wikilink = re.compile(f"^(?:{media_prefixes}):", flags=re.IGNORECASE | re.UNICODE)

    def rm_wikilink(obj):
        return bool(re_rm_wikilink.match(str(obj.title)))

    # Filters for references and tables
    def rm_tag(obj):
        return str(obj.tag) in {"ref", "table"}

    # Leave category links in-place but remove the category prefixes
    cat_prefixes = "|".join(["Category"] + ALIASES_CATEGORY.get(lang, []))
    re_clean_wikilink = re.compile(f"^(?:{cat_prefixes}):", flags=re.IGNORECASE | re.UNICODE)

    def is_category(obj):
        return bool(re_clean_wikilink.match(str(obj.title)))

    def clean_wikilink(obj):
        text = obj.__strip__()
        text = re.sub(re_clean_wikilink, "", text)
        obj.text = text

    def try_replace_obj(obj):
        try:
            clean_wikilink(obj)
        except ValueError:
            # For unknown reasons, objects are sometimes not found.
            pass

    def try_remove_obj(obj, section):
        try:
            section.remove(obj)
        except ValueError:
            # For unknown reasons, objects are sometimes not found.
            pass

    # Filter individual sections to clean.
    sections_text = []
    for section in wikicode.get_sections(flat=True, include_lead=True, include_headings=True):
        for obj in section.ifilter_wikilinks(recursive=True):
            if rm_wikilink(obj):
                try_remove_obj(obj, section)
            elif is_category(obj):
                try_replace_obj(obj)
        for obj in section.ifilter_tags(matches=rm_tag, recursive=True):
            try_remove_obj(obj, section)

        section_text = section.strip_code().strip()
        section_text = clean_magic(section_text)
        section_text = clean_box(section_text)
        section_text = clean_lsw(section_text)
        sections_text.append(section_text)

    return "\n\n".join(sections_text)


def map_func(sample, lang) -> bytes | None:
    namespace = sample.get('namespace')
    if namespace != '0':
        return None

    title = sample.get('title')
    if not title:
        return None
    title = str(title)

    content = sample.get('text')
    if not content:
        return None
    content = str(content)
    content = parse_and_clean_wikicode(content, lang)

    text = title + "\n\n" + content
    return text.encode("utf-8", errors="ignore")


def main(lang, date, *, num_workers=8):
    raw_dataset_dir_path = os.path.join(datasets_dir_path, "wikipedia", lang, date)
    raw_dataset = load_dataset(raw_dataset_dir_path)

    text_dataset_dir_path = os.path.join(datasets_dir_path, "wikipedia-text", lang, date)
    confirm_clear_dir_path(text_dataset_dir_path)
    text_dataset_writer = ChunkedBinaryDatasetWriter(text_dataset_dir_path)

    iterator = tqdm(raw_dataset)
    pool = multiprocessing.Pool(num_workers)
    for sample in pool.imap_unordered(partial(map_func, lang=lang), iterator):
        if sample is None:
            continue
        text_dataset_writer.append(sample)


if __name__ == '__main__':
    main(lang="zh", date="20231101", num_workers=16)
